//      gene_processing.cc
//
//      Basic sequence processing, read/write genome, annotation file implementation
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <iostream>
#include <boost/algorithm/string.hpp>
#include <boost/foreach.hpp>
#include <algorithm>
#include <set>
#include <armadillo>
#include "stl-helper.h"
#include "gene_processing.h"

/*********************** Genome **************************/

Genome::Genome()
  : __BASE()
{
}

Genome::Genome(const std::string& s)
  : __BASE(s)
{
}

void Genome::read_genome(const std::string& genomeFile)
{
  std::ifstream input(genomeFile.c_str());
  if (!input.good())
    FATAL_ERROR("Cannot open " << genomeFile << " to read");
 
  std::getline(input, firstLine);
  MSG(firstLine);
  int count = 0;
  while (!input.eof()) {
    count ++;
    std::string line;
    std::getline(input, line);

    // check last characters
    /*
    if (sz(line) > 0) {
      char c = line.at(sz(line)-1);    
      while ( (c != 'A' && c != 'T' && c != 'C' && c != 'G') && sz(line) > 0 ) {
	line.erase(sz(line)-1);
	c = line.at(sz(line)-1);
      }
    }
	*/
	std::string tmp;
	BOOST_FOREACH(char c, line)
	  if ( c == 'A' || c == 'T' || c == 'C' || c == 'G') tmp += c;
    (*this) += tmp;
    //MSG(line);
  }

  // verify
  for (int i = 0; i < sz(*this); i++) {
    char c = this->at(i);
    if (c != 'A' && c != 'T' && c != 'C' && c != 'G') 
      FATAL_ERROR("Check " << i << "-th position c = -" << (int)c << "-");
  }

  input.close();  
}

/*********************** Annotation **************************/

void _annotation::print() 
{
  MSG("id = " << gene_id << 
      " name = " << gene_name <<
      " number = " << gene_number <<
      " direction = " << gene_direction <<
      " start = " << gene_start <<
      " stop = " << gene_stop);
}

void Annotation::read_annotation_simple(const std::string& annotationFile)
{
  this->clear();

  std::ifstream input;
  
  FileAccess::open2read(input, annotationFile);

  int count = 0;
  while (!input.eof()) {
    std::string line = FileAccess::getline(input);
    //MSG(line);
    std::vector<std::string> tokens;
    boost::split(tokens, line, boost::is_any_of("\t \n"), boost::token_compress_on);
    if (sz(tokens) < 3) continue;

    _annotation a;
    char count_s[100];
    bool is_gene = false;
    count++;
    sprintf(count_s, "%d", count);

    if (tokens[0][0] == '#') {
      is_gene = true;
      tokens[0].erase(0,1);
    }

    a.gene_id = count_s;
    a.gene_name = std::string("orf-") + count_s + (is_gene?"#":"");
    a.gene_number = count;
    a.gene_direction = tokens[2];
    a.gene_start = RAND::s2type<int>(tokens[0]);
    a.gene_stop = RAND::s2type<int>(tokens[1]);
    //a.print();

    this->pb(a);
  }
  
  input.close();
}

void Annotation::read_annotation_simple1(const std::string& annotationFile)
{
  this->clear();

  std::ifstream input;
  
  FileAccess::open2read(input, annotationFile);

  int count = 0;
  while (!input.eof()) {
    std::string line = FileAccess::getline(input);
    //MSG(line);
    std::vector<std::string> tokens;
    boost::split(tokens, line, boost::is_any_of("\t \n"), boost::token_compress_on);
    if (sz(tokens) < 2) continue;

    _annotation a;
    char count_s[100];
    bool is_gene = false;
    count++;
    sprintf(count_s, "%d", count);

    if (tokens[0][0] == '#') {
      is_gene = true;
      tokens[0].erase(0,1);
    }

    a.gene_id = count_s;
    a.gene_name = std::string("orf-") + count_s + (is_gene?"#":"");
    a.gene_number = count;
    a.gene_direction = tokens[1] == "+" ? "Clockwise" : "CounterClockwise";
    
    MSG(RAND::s2type<int>(tokens[0]) << " " << a.gene_name << " " << a.gene_direction);
    if (a.gene_direction == "Clockwise") {
      a.gene_start = RAND::s2type<int>(tokens[0]) - 1; // convert 1 base to 0 base
      a.gene_stop = a.gene_start + 20;
    }
    else {
      a.gene_start = RAND::s2type<int>(tokens[0]) - 1 - 20;
      a.gene_stop = a.gene_start + 20;
    }
    //a.print();

    this->pb(a);
  }
  
  input.close();
}

void Annotation::read_annotation(const std::string& annotationFile)
{
  this->clear();

  std::ifstream input;
  
  FileAccess::open2read(input, annotationFile);

  while (!input.eof()) {
    std::string line = FileAccess::getline(input);
    //MSG(line);
    std::vector<std::string> tokens;
    boost::split(tokens, line, boost::is_any_of("\t \n"), boost::token_compress_on);
    if (sz(tokens) < 6) continue;

    _annotation a;
    a.gene_id = tokens[0];
    a.gene_name = tokens[1];
    a.gene_number = RAND::s2type<int>(tokens[2]);
    a.gene_direction = tokens[3];
    a.gene_start = RAND::s2type<int>(tokens[4]);
    a.gene_stop = RAND::s2type<int>(tokens[5]);
    //a.print();

    this->pb(a);
  }
  
  input.close();
}

std::string complementary_strand(const std::string& seq) {
  std::string tmp = seq;
  for (int i = 0; i < sz(tmp); i++) {
    char c = tmp.at(i);
    switch (c) {
    case 'A' : tmp.at(i) = 'T'; break;
    case 'T' : tmp.at(i) = 'A'; break;
    case 'C' : tmp.at(i) = 'G'; break;
    case 'G' : tmp.at(i) = 'C'; break;
    default: FATAL_ERROR("check i = " << i << " c = " << c);
    }      
  }
  return tmp;
}

void Annotation::get_gene_sequence(const Genome& genome, int n_upstream)
{
  for (int a_idx = 0; a_idx < sz(*this); a_idx++) {
    _annotation& a = this->at(a_idx);

    a.gene = genome.substr(a.gene_start, a.length());
    if ( a.gene_direction.compare("Clockwise") ) { // complementary strand
      std::reverse(a.gene.begin(), a.gene.end());      
      a.gene = complementary_strand(a.gene);

      a.gene_upstream = genome.substr(a.gene_stop+1, n_upstream);
      std::reverse(a.gene_upstream.begin(), a.gene_upstream.end());
      a.gene_upstream = complementary_strand(a.gene_upstream);

      //MSG(a.gene_direction);
    }
    else {
      a.gene_upstream = genome.substr(a.gene_start-n_upstream, n_upstream);
    }
    //MSG("gene=" << a.gene);
    //MSG(a.gene_upstream << " --- " << a.gene);
  }
}

void Annotation::save_upstream(const std::string& upstreamFile)
{
  std::ofstream output;
  FileAccess::open2write(output, upstreamFile);

  BOOST_FOREACH(const  _annotation& a, *this) {
    output << ">" << a.gene_id << " " << a.gene_name << " "
	   << a.gene_direction << " ";
    
    if ( a.gene_direction.compare("Clockwise") )
        output << a.gene_stop +1+1 << " " << a.gene_stop + sz(a.gene_upstream)+1;
    else
        output << a.gene_start - sz(a.gene_upstream)+1 << " " << a.gene_start-1+1;

    output  << " " << sz(a.gene_upstream) << " upstream: " << ENDL;
    output << a.gene_upstream << ENDL;
  }

  output.close();
}

void Annotation::all_upstream(Sequences& seqs)
{
  seqs.clear();
  BOOST_FOREACH(const _annotation& a, *this) {
    seqs.pb(a.gene_upstream);
  }

}

/*********************** Sequences **************************/

bool Sequences::compare(const std::string& pattern, const std::string& seq) 
{
  if (sz(pattern) != sz(seq))
    return false;
  
  for (int i = 0; i < sz(pattern); i++) {
    char p_c = pattern.at(i), s_c = seq.at(i);
    if (p_c != 'x' && p_c != s_c)
      return false;
  }
  return true;
}

void Sequences::check_occurence(const std::string& pattern, const std::string& seq, std::vector<int>& pos) 
{
  pos.clear();
  for (int i = 0; i <= sz(seq) - sz(pattern); i++) {
    if (compare(pattern, seq.substr(i, sz(pattern))))
      pos.pb(i);
  }
}

void Sequences::count_occurence(const std::string& pattern, std::vector<int>& count) const
{
  int n = max_length();
  count = std::vector<int>(n, 0);

  BOOST_FOREACH(const std::string& s, *this) {
    std::vector<int> pos;
    check_occurence(pattern, s, pos);
    BOOST_FOREACH(int p, pos) {
      count[p]++;
    }
  }
}

int Sequences::max_length() const
{
  int max = 0;
  BOOST_FOREACH(const std::string& s, *this) {
    if (sz(s) > max)
      max = sz(s);
  }
  return max;
}

void Sequences::read_sequence(const std::string& seqFile)
{
  this->clear();
  annotation_.clear();

  std::ifstream input;  
  FileAccess::open2read(input, seqFile);

  while (!input.eof()) {
    std::string line;
    std::string tmp;
    std::getline(input, line);
    if (sz(line) == 0) continue;
    if (line[0] == '>') { // annotation
      tmp = line;
      std::getline(input, line); // get next line for sequence
    }
    else					
      tmp = ">No information";

    // clear unwanted characters at the end of line
    if (sz(line) > 0) {
      char c = line.at(sz(line)-1);    
      while ( (c != 'A' && c != 'T' && c != 'C' && c != 'G') && sz(line) > 0 ) {
	line.erase(sz(line)-1);
	c = line.at(sz(line)-1);
      }
    }
    if (sz(line) > 0) {
      this->pb(line);
      annotation_.pb(tmp);
    }
  }

  input.close();
}

void GroupSequences::group_sequences(const Sequences& seqs)
{
  for (int i = 0; i < sz(seqs); i++) {
    const std::string& anno_i = seqs.annotation_[i];
    if (anno_i.find('#') != std::string::npos) {
      this->pb(Sequences());
    }
    this->back().pb(seqs[i]);
    this->back().annotation_.pb(anno_i);
  }
}

GroupSequences GroupSequences::removeSharp() const
{
  GroupSequences g;
  BOOST_FOREACH(const Sequences& s, *this) {
    //if (sz(g) == 0 || sz(g.back()) > 0) g.pb(Sequences());
    
    if (sz(s) == 1) 
      g.pb(s);
    else {
      g.pb(Sequences());
      for (int i = 0; i < sz(s); i++) {
        if (s.annotation_[i].find('#') == std::string::npos) {
          g.back().pb(s[i]);
          g.back().annotation_.pb(s.annotation_[i]);
        }
      }    
    }
  }
  //if (sz(g) > 0 && sz(g.back()) == 0) g.pop_back();
  return g;
}

/*********************** Pattern Generator **************************/
PatternGenerator::PatternGenerator(boost::program_options::variables_map& vm, const Sequences& seqs, const std::string& entropyFile)
  : vm_(vm), seqs_(seqs), entropyFile_(entropyFile)
{
}

void PatternGenerator::run()
{
  std::ofstream output(entropyFile_.c_str());

  std::string atcg = "ATCG";
  std::set<std::string> generated;
  for (int i = 0; i < 4; i++) {
    pattern_.pb(atcg.substr(i, 1));
    generated.insert(atcg.substr(i, 1));
  }

  int begin = 0, end = 4;
  int max_length = 1;
  int motif_length = vm_["mtlen"].as<int>();

  output << "id" << "\t " << "pattern" << "\t " << "total_count" << "\t " << "entropy" << "\t " << "max_entropy"
	 << "\t " << "entropy/max_entropy";
  for (int i = -seqs_.max_length(); i < 0; i++)
    output << "\t pos(" << i << ")";
  output << ENDL;
  
  int index = 1;
  while (begin < end) {
    std::string pattern = pattern_[begin++];
    if (sz(pattern) > max_length) {
      max_length = sz(pattern);
      MSG("max-length = " << max_length << " begin = " << begin);
    }

    /** Count and compute entropy of position frequency of the pattern */
    vi_t count;
    seqs_.count_occurence(pattern, count);
    double entropy = RAND::get_entropy(all(count));
    double non_zero = 0;
    BOOST_FOREACH(int c, count)
      if (c > 0) non_zero += 1;
    double max_entropy = log(non_zero);
    int total_count = std::accumulate(all(count), 0);
    
    //if (total_count < sz(seqs_) * 0.05) continue;

    output << index++ << "\t " << pattern << "\t " << total_count << "\t " << entropy 
	   << "\t " << max_entropy << "\t " << entropy/max_entropy;
    BOOST_FOREACH(int c, count) 
      output << "\t " << c;
    output << ENDL;

    /** compute next line of candidates */
    //double threshold[] = {1.0, 0.001, 0.005, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01};
    //if ((max_entropy - entropy) / max_entropy < threshold[sz(pattern)]) continue;
    std::vector<std::string> next = next_candidate(pattern);
    BOOST_FOREACH(const std::string p, next) {
      if (sz(p) <= motif_length && generated.find(p) == generated.end()) {
	pattern_.pb(p);
	generated.insert(p);
	end++;
      }
    }
  }

  output.close();
}

std::vector<std::string> PatternGenerator::next_candidate(const std::string& pattern)
{
  std::vector<std::string> next;
  std::string atcg = "ATCG";
  for (int i = 0; i < 4; i++)
    next.pb(pattern + atcg[i]);
  if (sz(pattern) >= 2) {
    for (int i = 1; i < sz(pattern); i++) {
      next.pb(pattern.substr(0, i) + "x" + pattern.substr(i));
    }
  }
  return next;
}

